from kubernetes.stream import stream
from langchain.agents import AgentExecutor, create_react_agent
# from langchain.llms import Ollama
from langchain.prompts import PromptTemplate
from langchain_core.callbacks import StreamingStdOutCallbackHandler
from langchain_ollama import OllamaLLM
from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings
import os


class DocumentStorage:
    def __init__(self, faiss_index_path="faiss_index/index.faiss"):
        self.faiss_index_path = faiss_index_path
        self.embeddings = HuggingFaceEmbeddings(
            model_name="model/bge-large-zh",
            model_kwargs={'device': 'cuda'}   # 如果没有GPU改成 'cpu'
        )
        self.vector_store = None

    def load_faiss_vector_store(self):
        if os.path.exists(self.faiss_index_path):
            print(f"加载本地FAISS向量存储: {self.faiss_index_path}")
            self.vector_store = FAISS.load_local(
                self.faiss_index_path,
                self.embeddings,
                allow_dangerous_deserialization=True
            )
        else:
            print(f"未检测到向量存储文件 {self.faiss_index_path}")

    def search_similar_documents(self, query, top_k=6):
        if not self.vector_store:
            print("❌ FAISS 向量存储未加载，请先创建或加载存储")
            return []
        similar_docs = self.vector_store.similarity_search(query, k=top_k)
        print(f"查询 {query} 的相关内容:\n")
        for i, doc in enumerate(similar_docs):
            print(f"文档 {i+1}:\n{doc.page_content[:200]}\n")
        return [doc.page_content for doc in similar_docs]



class ReactAgentSystem:
    def __init__(self, document_storage):
        self.document_storage = document_storage
        self.llm = OllamaLLM(
            model="qwen3:4b",
            temperature=0.3,
            num_ctx=4096,
            stream = True,
            callbacks=[StreamingStdOutCallbackHandler()]
        )

    # def create_react_agent(self, query, related_docs):
    #     react_prompt = PromptTemplate(
    #         template=(
    #             "以下是与问题相关的企业文档内容：\n{docs}\n\n"
    #             "请基于文档信息回答问题：\n问题：{query}\n"
    #             # "如果在文档中没有找到相关内容就直接告诉我没有"
    #         ),
    #         input_variables=["query", "related_docs"]
    #     )
    #     return react_prompt.format(query=query, docs="\n".join(related_docs))
    def create_react_agent(self, query, related_docs):
        react_prompt = PromptTemplate(
            template=(
                """作为一个精确的RAG系统助手，请严格按照以下指南回答用户问题：
                    1. 仔细分析问题，识别关键词和核心概念。
                    2. 从提供的上下文中精确定位相关信息，优先使用完全匹配的内容。
                    3. 构建回答时，确保包含所有必要的关键词，提高关键词评分(scoreikw)。
                    4. 保持回答与原文的语义相似度，以提高向量相似度评分(scoreies)。
                    5. 控制回答长度，理想情况下不超过参考上下文长度的1.5倍，最多不超过2.5倍。
                    6. 对于表格查询或需要多段落/多文档综合的问题，给予特别关注并提供更全面的回答。
                    7. 如果上下文信息不足，可以进行合理推理，但要明确指出推理部分。
                    8. 回答应简洁、准确、完整，直接解答问题，避免不必要的解释。
                    9. 不要输出“检索到的文本块”、“根据”，“信息”等前缀修饰句，直接输出答案即可
                    10. 不要使用"根据提供的信息"、"支撑信息显示"等前缀，直接给出答案。
                    问题: {query}
                    支撑信息：
                    ···
                    {docs}
                    ···
                    请提供准确、相关且简洁的回答："""
                # "如果在文档中没有找到相关内容就直接告诉我没有"
            ),
            input_variables=["query", "related_docs"]
        )
        return react_prompt.format(query=query, docs="\n".join(related_docs))

    def generate_response(self, query):
        related_docs = self.document_storage.search_similar_documents(query)
        print(f"相关文档{related_docs}")
        prompt = self.create_react_agent(query, related_docs)
        response = self.llm.invoke(prompt)
        print("\nReact Agent回答：\n", response)
        return response


class ZeroShotAgentSystem:
    def __init__(self):
        self.llm = OllamaLLM(
            model="qwen3:4b",
            temperature=0.3,
            num_ctx=4096,
            stream = True,
            callbacks=[StreamingStdOutCallbackHandler()]
        )

    def create_zero_shot_prompt(self, query):
        prompt = PromptTemplate(
            template="请基于常识回答下述问题：\n问题：{query}\n",
            input_variables=["query"]
        )
        return prompt.format(query=query)

    def generate_response(self, query):
        prompt = self.create_zero_shot_prompt(query)
        response = self.llm.invoke(prompt)   # ✅ 修正
        print("\nZero-shot Agent回答：\n", response)
        return response


class EnterpriseAgentSystem:
    def __init__(self, faiss_index_path):
        self.document_storage = DocumentStorage(faiss_index_path)
        self.document_storage.load_faiss_vector_store()  # ⬅️ 加上这行
        self.react_agent = ReactAgentSystem(self.document_storage)
        self.zero_shot_agent = ZeroShotAgentSystem()

    def execute_query(self, query, user_react=True):
        related_docs = self.document_storage.search_similar_documents(query)
        if user_react:
            print("使用ReAct Agent进行回答：")
            # 直接调用 react_agent 来生成回答
            response = self.react_agent.generate_response(query)
        else:
            print("使用Zero-shot Agent进行回答：")
            response = self.zero_shot_agent.generate_response(query)

        print(f"最终回答：{response}")
        return response



if __name__ == "__main__":
    # documents_directory = "data/1.txt"
    # enterprise_agent = EnterpriseAgentSystem(documents_directory)
    # query1 = "RK3588淘宝链接"
    # print("查询1结果")
    # enterprise_agent.execute_query(query1, user_react=True)
    #
    # query2 = "申请部门申请日期"
    # print("查询2结果")
    # enterprise_agent.execute_query(query2, user_react=True)
    if __name__ == "__main__":
        # 初始化企业代理系统
        faiss_index_path = "faiss_index"

        enterprise_agent = EnterpriseAgentSystem(faiss_index_path)

        print("欢迎使用企业问答系统！请输入您的问题（输入'退出'结束对话）")

        # 进入问答循环
        while True:
            # 获取用户输入
            user_query = input("\n请输入您的问题: ")

            # 检查是否退出
            if user_query.strip().lower() == "退出":
                print("感谢使用，再见！")
                break

            # 执行查询并获取结果
            print("查询结果:")
            enterprise_agent.execute_query(user_query, user_react=True)